from time import time

import matplotlib.pyplot as plt
import numpy as np
from sklearn import (manifold, datasets)


# %%
# 将3维数据可视化
def plot_embedding_3d(X, y, title=None):
    # 坐标缩放到[0,1]区间
    x_min, x_max = np.min(X, axis=0), np.max(X, axis=0)
    X = (X - x_min) / (x_max - x_min)
    # 降维后的坐标为（X[i, 0], X[i, 1],X[i,2]），在该位置画出对应的digits
    fig = plt.figure()
    ax = fig.add_subplot(1, 1, 1, projection='3d')
    for i in range(X.shape[0]):
        ax.text(X[i, 0], X[i, 1], X[i, 2], str(y[i])
                , color=plt.cm.Set1(y[i] / 10.)
                , fontdict={'weight': 'bold', 'size': 9})
        if title is not None:
            plt.title(title)
    plt.show()


# %%
# 将降维后的数据可视化,2维
def plot_embedding_2d(X, y, title=None):
    # 坐标缩放到[0,1]区间
    x_min, x_max = np.min(X, axis=0), np.max(X, axis=0)
    X = (X - x_min) / (x_max - x_min)

    # 降维后的坐标为（X[i, 0], X[i, 1]），在该位置画出对应的digits
    fig = plt.figure()
    ax = fig.add_subplot(1, 1, 1)
    for i in range(X.shape[0]):
        ax.text(X[i, 0], X[i, 1], str(y[i]),
                color=plt.cm.Set1(y[i] / 10.),
                fontdict={'weight': 'bold', 'size': 9})

    if title is not None:
        plt.title(title)
    plt.show()


# %%  加载数据

digits = datasets.load_digits(n_class=5)
X = digits.data
y = digits.target
print(X.shape)
n_img_per_row = 20
img = np.zeros((10 * n_img_per_row, 10 * n_img_per_row))
for i in range(n_img_per_row):
    ix = 10 * i + 1
    for j in range(n_img_per_row):
        iy = 10 * j + 1
        img[ix:ix + 8, iy:iy + 8] = X[i * n_img_per_row + j].reshape((8, 8))
        plt.imshow(img, cmap=plt.cm.binary)
        plt.title('A selection from the 64-dimensional digits dataset')

print("Computing t-SNE embedding")
tsne = manifold.TSNE(n_components=3, init='pca', random_state=0)
t0 = time()
X_tsne = tsne.fit_transform(X)
plot_embedding_2d(X_tsne[:, 0:2], y, "t-SNE 2D")
plot_embedding_3d(X_tsne, y, "t-SNE 3D (time %.2fs)" % (time() - t0))
